﻿#include "capture.h"
#include <map>
#include <chrono>
#include <thread>
using namespace std;
vector<string> user_pass;  // 用于存储FTP用户和密码，当两个字符串都获取后进行处理
int counter;  // 分钟计数
vector<string> from_ip;  // 存储源IP
vector<string> to_ip;    // 存储目的IP
vector<long long> from_ip_data;  // 存储源IP发送的数据量
vector<long long> to_ip_data;    // 存储目的IP接收的数据量
vector<string> from_mac;  // 存储源MAC地址
vector<string> to_mac;    // 存储目标MAC地址
vector<long long> from_mac_data;  // 存储源MAC发送的数据量
vector<long long> to_mac_data;    // 存储目的MAC接收的数据量

// 该函数处理捕获到的MAC和IP数据包
void mac_ip_packet_handler(unsigned char* param, const struct pcap_pkthdr* header, const unsigned char* pkt_data)
{
	// 转换时间格式
	string date_time;  // 存储格式化后的时间字符串
	char timestr[16];  // 缓存区，存储时间字符串
	time_t local_tv_sec = header->ts.tv_sec;  // 从捕获头中获取数据包的秒级时间戳
	tm* ltime = localtime(&local_tv_sec);  // 将时间戳转换为本地时间
	strftime(timestr, sizeof timestr, "%H:%M:%S", ltime);  // 改变时间格式
	// 格式化完整的时间（年-月-日 小时:分钟:秒）
	date_time = to_string(ltime->tm_year + 1900) + '-' +
		to_string(ltime->tm_mon + 1) + '-' +
		to_string(ltime->tm_mday) + ' ' + timestr;

	//检索数据报头部的位置
	mac_header* mh;  // 以太网帧头
	ip_header* ih;   // IP包头
	mh = (mac_header*)pkt_data;  // pkt_data转换为以太网帧头
	ih = (ip_header*)(pkt_data + 14);  // IP头偏移14字节

	//解析并格式化MAC地址
	stringstream ss;  // 拼接MAC地址
	string mac_dst;   // 目标MAC地址
	string mac_src;   // 源MAC地址
	// 解析目标MAC地址
	for (int i = 0; i < 6; ++i) {  // 以太网地址长度为6字节
		ss << hex << setw(2) << setfill('0') << (int)mh->dst_addr[i];  // 字节转换为十六进制字符串
		mac_dst += ss.str();  // 拼接MAC地址
		if (i < 5) {
			mac_dst += '-';  // 改变MAC地址的格式
		}
		ss.str("");  // 清空字符串流
	}
	// 解析源MAC地址
	for (int i = 0; i < 6; ++i) {
		ss << hex << setw(2) << setfill('0') << (int)mh->src_addr[i];  // 同理处理源MAC地址
		mac_src += ss.str();
		if (i < 5) {
			mac_src += '-';
		}
		ss.str("");
	}

	//解析并格式化IP地址
	string ip_dst;  // 目标IP地址
	string ip_src;  // 源IP地址
	// 解析目标IP地址
	for (int i = 0; i < 4; i++) {  // IPv4地址长度为4字节
		ip_dst += to_string(ih->daddr.byte[i]);  // 将字节转换为十进制字符串
		if (i < 3) {
			ip_dst += '.';  // 改变IP地址格式
		}
	}
	// 解析源IP地址，同理
	for (int i = 0; i < 4; i++) {
		ip_src += to_string(ih->saddr.byte[i]); 
		if (i < 3) {
			ip_src += '.';
		}
	}

	// 输出捕获到的包信息：时间、源MAC、源IP、目标MAC、目标IP、数据包长度
	cout << date_time + ", " + mac_src + ", " + ip_src + ", " + mac_dst + ", " + ip_dst + ", " 
		<< header->len << "B" << endl;

	// 将信息追加写入到MAC_IP.csv文件中
	ofstream outFile("MAC_IP.csv", ios::app);
	outFile << date_time << ',' << mac_src << ',' << ip_src << ',' << mac_dst << ',' << ip_dst << ',' << header->len << ',' << endl;
	outFile.close();

	//每分钟统计来自不同 IP 地址和 MAC 地址的数据长度
	int i;

	// 统计源MAC地址发送的数据量
	for (i = 0; i < from_mac.size(); i++) {
		if (from_mac[i] == mac_src) {  // 如果找到了匹配的源MAC
			from_mac_data[i] += header->len;  // 累加该MAC地址发送的数据长度
			break;
		}
	}
	if (i == from_mac.size()) {  // 如果没有找到该源MAC，则添加新的MAC和数据量记录
		from_mac.push_back(mac_src);
		from_mac_data.push_back(header->len);
	}

	// 统计目标MAC地址接收的数据量
	for (i = 0; i < to_mac.size(); i++) {
		if (to_mac[i] == mac_dst) {  // 如果找到了匹配的目标MAC
			to_mac_data[i] += header->len;  // 累加该MAC地址接收的数据长度
			break;
		}
	}
	if (i == to_mac.size()) {  // 如果没有找到该目标MAC，则添加新的MAC和数据量记录
		to_mac.push_back(mac_dst);
		to_mac_data.push_back(header->len);
	}

	// 统计源IP地址发送的数据量
	for (i = 0; i < from_ip.size(); i++) {
		if (from_ip[i] == ip_src) {  // 如果找到了匹配的源IP
			from_ip_data[i] += ih->tlen;  // 累加该IP地址发送的数据长度
			break;
		}
	}
	if (i == from_ip.size()) {  // 如果没有找到该源IP，则添加新的IP和数据量记录
		from_ip.push_back(ip_src);
		from_ip_data.push_back(ih->tlen);
	}

	// 统计目标IP地址接收的数据量
	for (i = 0; i < to_ip.size(); i++) {
		if (to_ip[i] == ip_dst) {  // 如果找到了匹配的目标IP
			to_ip_data[i] += ih->tlen;  // 累加该IP地址接收的数据长度
			break;
		}
	}
	if (i == to_ip.size()) {  // 如果没有找到该目标IP，则添加新的IP和数据量记录
		to_ip.push_back(ip_dst);
		to_ip_data.push_back(ih->tlen);
	}

	// 检查是否到达了新的分钟
	if (counter != ltime->tm_min) {  // 如果当前分钟与上次记录的分钟不同
		counter = ltime->tm_min;  // 更新计数器为当前分钟
		cout << endl;
		ofstream coutFile("MAC_IP_MinRecord.csv", ios::app);
		// 输出过去一分钟内每个源MAC发送的数据长度
		for (int j = 0; j < from_mac.size(); j++) {
			cout << "The length of data sent by MAC:" + from_mac[j] + " in the last minute: " << from_mac_data[j] << endl;
			coutFile << "The length of data sent by MAC:" + from_mac[j] + " in the last minute: " << from_mac_data[j] << endl;
		}
		cout << endl;
		coutFile << endl;
		// 输出过去一分钟内每个目标MAC接收的数据长度
		for (int j = 0; j < to_mac.size(); j++) {
			cout << "The length of data received on MAC:" + to_mac[j] + " in the last minute: " << to_mac_data[j] << endl;
			coutFile << "The length of data received on MAC:" + to_mac[j] + " in the last minute: " << to_mac_data[j] << endl;
		}
		cout << endl;

		// 输出过去一分钟内每个源IP发送的数据长度
		for (int j = 0; j < from_ip.size(); j++) {
			cout << "The length of data sent by IP:" + from_ip[j] + " in the last minute: " << from_ip_data[j] << endl;
			coutFile << "The length of data sent by IP:" + from_ip[j] + " in the last minute: " << from_ip_data[j] << endl;
		}
		cout << endl;

		// 输出过去一分钟内每个目标IP接收的数据长度
		for (int j = 0; j < to_ip.size(); j++) {
			cout << "The length of data received on IP:" + to_ip[j] + " in the last minute: " << to_ip_data[j] << endl;
			coutFile << "The length of data received on IP:" + to_ip[j] + " in the last minute: " << to_ip_data[j] << endl;
		}
		cout << endl;
		coutFile << endl;
		coutFile.close();  // 关闭文件
		// 清空MAC和IP数据量记录，以便统计下一个时间段的数据
		from_mac.clear();
		to_mac.clear();
		from_mac_data.clear();
		to_mac_data.clear();
		from_ip.clear();
		to_ip.clear();
		from_ip_data.clear();
		to_ip_data.clear();
	}
}


// 该函数处理捕获到的FTP数据包
void ftp_packet_handler(unsigned char* param, const struct pcap_pkthdr* header, const unsigned char* pkt_data)
{
	// 转换时间格式，同上
	string date_time;
	char timestr[16]; 
	time_t local_tv_sec = header->ts.tv_sec;
	tm* ltime = localtime(&local_tv_sec);  
	strftime(timestr, sizeof timestr, "%H:%M:%S", ltime); 
	date_time = to_string(ltime->tm_year + 1900) + '-' +
		to_string(ltime->tm_mon + 1) + '-' +
		to_string(ltime->tm_mday) + ' ' + timestr;

	//检索数据包头部的位置
	mac_header* mh;  // 以太网帧头部
	ip_header* ih;   // IP头部
	mh = (mac_header*)pkt_data;  // 以太网帧头从pkt_data的起始位置开始
	ih = (ip_header*)(pkt_data + 14);  // IP头偏移14字节

	//解析并格式化MAC地址
	stringstream ss;  // 字符串流用于格式化MAC地址
	string mac_dst;   // 目标MAC地址
	string mac_src;   // 源MAC地址
	// 解析目标MAC地址（6字节）
	for (int i = 0; i < 6; i++) {
		ss << hex << setw(2) << setfill('0') << (int)mh->dst_addr[i];  // 转换为十六进制字符串
		mac_dst += ss.str();
		if (i < 5) mac_dst += '-';  // 每两个字节之间用'-'分隔
		ss.str("");  // 清空字符串流
	}
	// 解析源MAC地址（6字节）
	for (int i = 0; i < 6; i++) {
		ss << hex << setw(2) << setfill('0') << (int)mh->src_addr[i];  // 同理处理源MAC地址
		mac_src += ss.str();
		if (i < 5) mac_src += '-';
		ss.str("");
	}

	//解析并格式化IP地址
	string ip_dst;  // 目标IP地址
	string ip_src;  // 源IP地址
	// 解析目标IP地址（4字节）
	for (int i = 0; i < 4; i++) {
		ip_dst += to_string(ih->daddr.byte[i]);  // 将每个字节转换为十进制
		if (i < 3) ip_dst += '.';  // 使用点号分隔IP地址的四部分
	}
	// 解析源IP地址（4字节）
	for (int i = 0; i < 4; i++) {
		ip_src += to_string(ih->saddr.byte[i]);  // 同理处理源IP地址
		if (i < 3) ip_src += '.';
	}

	//提取FTP数据
	int ip_header_len = (ih->ver_ihl & 0x0F) * 4;  // IP头长度
	int tcp_header_len = ((unsigned char)pkt_data[14 + ip_header_len + 12] >> 4) * 4;  // TCP头长度
	string ftp_data = (char*)(pkt_data + 14 + ip_header_len + tcp_header_len);  // FTP数据部分
	//cout << ftp_data << endl;  // 用于调试，查看完整的FTP数据

	int idx;  // 用于存储字符串查找的索引
	// user_pass列表暂时为空且找到了"USER"，提取用户名
	if (user_pass.size() == 0 && (idx = ftp_data.find("USER", 0)) != -1) {
		int index = ftp_data.find('\r', idx + 4);  // 查找回车符（命令结束的位置）
		user_pass.push_back(ftp_data.substr(idx + 5, index - (idx + 5)));  // 提取用户名并存储到user_pass中
	}
	// user_pass已经有一个元素且找到了"PASS"命令，提取密码
	else if (user_pass.size() == 1 && (idx = ftp_data.find("PASS", 0)) != -1) {
		int index = ftp_data.find('\r', idx + 4);  // 查找回车符（命令结束的位置）
		user_pass.push_back(ftp_data.substr(idx + 5, index - (idx + 5)));  // 提取密码并存储到user_pass中
	}
	// 当用户名和密码都已获取（user_pass有两个元素时），处理登录结果
	else if (user_pass.size() == 2) {
		// FTP响应代码230表示登录成功
		if (ftp_data.find("230", 0) != -1) {
			//输出成功登录信息
			cout << date_time + ", " + mac_src + ", " + ip_src + ", " + mac_dst + ", " + ip_dst + ", "
				<< user_pass[0] << ", " << user_pass[1] << ", " << "SUCCEED" << endl;

			// 将信息写入FTP.csv文件
			ofstream outFile("FTP.csv", ios::app);
			outFile << date_time << ',' << mac_src << ',' << ip_src << ',' << mac_dst << ',' << ip_dst << ','
				<< user_pass[0] << ", " << user_pass[1] << ", " << "SUCCEED" << endl;
			outFile.close();

			// 清空user_pass以便处理下一个FTP会话
			user_pass.clear();
		}
		// 如果FTP响应代码为530，表示登录失败
		else if (ftp_data.find("530", 0) != -1) {
			/* 按要求格式输出失败登录信息 */
			cout << date_time + ", " + mac_src + ", " + ip_src + ", " + mac_dst + ", " + ip_dst + ", "
				<< user_pass[0] << ", " << user_pass[1] << ", " << "FAILED" << endl;

			// 将失败登录信息写入FTP.csv文件
			ofstream outFile("FTP.csv", ios::app);
			outFile << date_time << ',' << mac_src << ',' << ip_src << ',' << mac_dst << ',' << ip_dst << ','
				<< user_pass[0] << ", " << user_pass[1] << ", " << "FAILED" << endl;
			outFile.close();

			// 清空user_pass以便处理下一个FTP会话
			user_pass.clear();
		}
	}
}

int main()
{
	bool TYPE;
	int tmp = 0;
	while (tmp != 1 && tmp != 2) {  
		cout << "Please select the task you want to perform:" << endl;
		cout << "1. Capture MAC and IP information" << endl;
		cout << "2. Capture FTP data" << endl;
		cin >> tmp;
	}
	TYPE = (tmp == 1);  // 根据用户输入选择任务类型，1表示捕获MAC和IP信息，2表示捕获FTP数据
	pcap_if_t* alldevs;  // 所有可用网络设备的链表头
	pcap_if_t* d;        // 当前选择的网络设备
	char errbuf[PCAP_ERRBUF_SIZE];  // 存储错误信息的缓冲区

	//检索设备列表
	if (pcap_findalldevs(&alldevs, errbuf) == -1)  // 调用pcap_findalldevs函数来获取所有网络设备列表
	{
		fprintf(stderr, "Error in pcap_findalldevs: %s\n", errbuf);  // 如果出错，输出错误信息
		exit(1);  // 程序退出
	}

	//输出设备列表
	int i = 0;  // 设备编号
	for (d = alldevs; d; d = d->next)  // 遍历设备列表，输出每个设备的信息
	{
		printf("%d. %s", ++i, d->name);  // 输出设备编号和设备名称
		if (d->description)  // 如果设备有描述，则输出描述
			printf(" (%s)\n", d->description);
		else
			printf(" (No description available)\n");  // 如果没有描述，输出“无描述信息”
	}

	if (i == 0)  // 如果没有找到任何设备
	{
		printf("\nNo interfaces found! Make sure WinPcap is installed.\n");
		return -1;  // 退出程序
	}

	int inum;  // 用户输入的设备编号
	printf("Enter the interface number (1-%d):", i);  // 提示用户选择一个设备
	cin >> inum;  // 获取用户输入

	//检查用户输入的设备编号是否有效
	if (inum < 1 || inum > i)  // 如果设备编号超出范围
	{
		printf("\nAdapter number out of range.\n");
		pcap_freealldevs(alldevs);  // 释放设备列表的内存
		return -1;  // 退出程序
	}

	// 根据用户输入的编号，定位到选中的设备
	for (d = alldevs, i = 1; i < inum; d = d->next, i++);  

	//启动适配器
	pcap_t* adhandle;  // 用于捕获数据包的句柄
	if ((adhandle = pcap_open_live(d->name, 65536, 1, 1000, errbuf)) == NULL)  // 打开指定的网络适配器
	{
		fprintf(stderr, "\nUnable to open the adapter. %s is not supported by WinPcap\n", d->name);  // 如果出错，输出错误信息
		pcap_freealldevs(alldevs);  // 释放设备列表的内存
		return -1;  // 退出
	}

	//检查数据链路层类型，只支持以太网
	if (pcap_datalink(adhandle) != DLT_EN10MB)  // 检查适配器的链路层类型是否为以太网
	{
		fprintf(stderr, "\nThis program works only on Ethernet networks.\n");
		pcap_freealldevs(alldevs);  // 释放设备列表的内存
		return -1;  // 退出程序
	}

	//获取网络掩码
	unsigned int netmask;
	if (d->addresses != NULL) {  // 如果设备有地址信息，获取网络掩码
		netmask = ((struct sockaddr_in*)(d->addresses->netmask))->sin_addr.S_un.S_addr;
	}
	else {
		netmask = 0xffffff;  // 如果没有地址信息，默认使用掩码0xffffff
	}

	//设置数据包过滤器
	string packet_filter;  // 过滤器字符串
	if (TYPE) {  // 如果TYPE为true，表示执行捕获IP、MAC数据包任务
		packet_filter = "ip";  // 过滤IP包
	}
	else {  // 否则执行捕获FTP数据包任务
		packet_filter = "tcp port 21";  // 过滤FTP数据包（端口21）
	}
	struct bpf_program fcode;  // 用于存储编译后的过滤规则
	if (pcap_compile(adhandle, &fcode, packet_filter.c_str(), 1, netmask) < 0)  // 编译过滤器表达式
	{
		fprintf(stderr, "\nUnable to compile the packet filter. Check the syntax.\n");
		pcap_freealldevs(alldevs);  // 释放设备列表的内存
		return -1;  // 退出程序
	}
	if (pcap_setfilter(adhandle, &fcode) < 0)  // 设置过滤器
	{
		fprintf(stderr, "\nError setting the filter.\n");
		pcap_freealldevs(alldevs);  // 释放设备列表的内存
		return -1;  // 退出程序
	}

	printf("\nlistening on %s...\n", d->description);  // 输出当前正在监听的设备信息
	pcap_freealldevs(alldevs);  // 释放设备列表的内存

	//开始捕获数据包并调用相应的回调函数
	if (TYPE) {  // 如果TYPE为true，执行MAC和IP信息捕获任务
		ofstream outFile("MAC_IP.csv", ios::out);  // 打开一个CSV文件用于记录MAC和IP信息
		outFile << "时间" << ',' << "源MAC" << ',' << "源IP" << ',' << "目标MAC" << ',' << "目标IP" << ',' << "帧长度" << ',' << endl;  // 写入CSV文件的表头
		outFile.close();
		pcap_loop(adhandle, 0, mac_ip_packet_handler, NULL);  // 开始捕获数据包，并使用mac_ip_packet_handler作为回调函数
	}
	else {  // 否则，执行FTP数据包捕获任务
		ofstream outFile("FTP.csv", ios::out);  // 打开一个CSV文件用于记录FTP信息
		outFile << "时间" << ',' << "源MAC" << ',' << "源IP" << ','
			<< "目标MAC" << ',' << "目标IP" << ','
			<< "登录名" << ',' << "口令" << ',' << "成功与否" << endl;  // 写入CSV文件的表头
		outFile.close();
		pcap_loop(adhandle, 0, ftp_packet_handler, NULL);  // 开始捕获数据包，并使用ftp_packet_handler作为回调函数
	}

	return 0;
}
